package com.lenovo.spider.component.downloader;

import cn.hutool.core.util.RandomUtil;
import cn.hutool.http.Header;
import cn.hutool.http.HttpRequest;
import cn.hutool.http.HttpResponse;
import com.lenovo.spider.ConfigHolder;
import com.lenovo.spider.common.spi.ServiceProvider;
import com.lenovo.spider.component.TaskProcess.Downloader;
import com.lenovo.spider.domain.Request;
import com.lenovo.spider.domain.Response;
import com.lenovo.spider.domain.Task;
import com.lenovo.spider.exception.SpiderException;
import lombok.extern.slf4j.Slf4j;
import lombok.val;

import java.net.HttpCookie;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicInteger;

import static com.lenovo.spider.Const.DOWNLOADER_USER_AGENT;
import static com.lenovo.spider.Const.MAX_HTTP_DOWNLOADER_COUNT;

/**
 * Http 下载器
 *
 * @author <a href="mailto:liuhan3@lenovo.com">HanL(liuhan3)</a>
 * on 2018/4/13 12:11
 */
@Slf4j
@ServiceProvider
public class HutoolHttpDownloader implements Downloader {

    private final AtomicInteger counter;
    private final int limitCount;

    public HutoolHttpDownloader() {
        this.counter = new AtomicInteger(0);
        this.limitCount = ConfigHolder.getInt(MAX_HTTP_DOWNLOADER_COUNT, 500);
    }

    @Override
    public boolean support(Task task) {
        return task.getRequest().getDownloaderType() == Type.HTTP;
    }

    @Override
    public Task handle(Task task) {
        int count = counter.incrementAndGet();
        if (count > limitCount) {
            throw new SpiderException.NeedRetryExcpetion("Http downloader count is more than {}", limitCount);
        }
        val request = task.getRequest();
        val response = createRequest(request).execute();
        return task.setResponse(convertResponse(request, response));
    }

    /**
     * 创建实际的 Http 请求
     *
     * @param request
     * @return
     */
    private HttpRequest createRequest(Request request) {
        val req = HttpRequest.get(request.getUrl());
        Optional.ofNullable(request.getProxy()).ifPresent(req::setProxy);
        Optional.ofNullable(request.getHeaders()).ifPresent(req::addHeaders);
        req.header(Header.USER_AGENT, RandomUtil.randomEle(ConfigHolder.getStrings(DOWNLOADER_USER_AGENT)));

        if (request.isUseCookie()) {
            request.getCookies().forEach((k, v) -> req.cookie(new HttpCookie(k, v)));
        } else {
            req.disableCookie();
        }
        req.setFollowRedirects(true);

        return req;
    }

    /**
     * 转换 Http 响应
     *
     * @param request
     * @param response
     * @return
     */
    private Response convertResponse(Request request, HttpResponse response) {
        val resp = new Response(response.getStatus(), request.getUrl(), request);
        resp.setCharset(response.charset());
        resp.setContent(response.body());
        return resp;
    }
}
